# 获取数据库会话和记忆管理器
# 执行向量相似度搜索
import asyncio
from datetime import datetime
from typing import List, Dict, Any

from langchain_core.tools import tool

from config.database import get_db_session
from memory.MemoryFactory import MemoryFactory
# 导入通用工具管理器
from tools.ToolManager import ToolManager, SingleToolConfig
import json


@tool
def get_memory_history(
        scene_id: str = "",
        agent_id: str = "",
        session_id: str = "",
        user_id: str = "1",
        keyword: str = "",
        tag: str = "",
        limit: int = 10,
        memory_type: str = "all"
) -> str:
    """
    从向量库ChromaDB中获取相关的历史记忆信息（基于语义相似度）。

    参数:
        scene_id: 可选，场景ID，用于筛选特定场景的记忆
        agent_id: 可选，智能体ID，用于筛选特定Agent的记忆
        session_id: 可选，会话ID，用于筛选特定对话的记忆
        user_id: 可选，用户ID，默认为"1"
        keyword: 必需，关键词，用于向量搜索相似内容的记忆
        tag: 可选，标签，用于筛选特定标签的记忆
        limit: 可选，返回记忆数量限制，默认10条
        memory_type: 可选，记忆类型。可选值：
            - "all": 返回所有类型的记忆（默认）
            - "conversation": 仅返回对话记忆
            - "important": 仅返回重要记忆
            - "recent": 仅返回最近记忆
    返回:
        相关历史记忆信息的JSON字符串（按相似度排序）。
    """
    try:
        # 检查关键词
        if not keyword or not keyword.strip():
            return json.dumps({
                "error": "需要提供关键词进行向量搜索",
                "suggestion": "请提供要搜索的关键词，系统将从向量库中找到语义相似的记忆"
            }, ensure_ascii=False)

        db_session = next(get_db_session())
        if not db_session:
            return json.dumps({"error": "无法连接数据库"}, ensure_ascii=False)

        try:
            # 创建记忆管理器，使用Mem0向量搜索
            memory_manager = MemoryFactory.create_memory_system(
                memory_config={"type": "mem0"},
                db_session=db_session,
                scene_id=scene_id or "search_context",
                agent_id=agent_id or "search_agent",
                session_id=session_id or "search_session",
                user_id=user_id
            )

            if not memory_manager or not memory_manager.mem0_provider:
                return json.dumps({"error": "无法创建向量搜索记忆管理器"}, ensure_ascii=False)

            # 执行同步搜索
            search_results = memory_manager.mem0_provider.mem0.search_memories(
                query=keyword,
                user_id=user_id,
                agent_id=agent_id if agent_id else None,
                scene_id=scene_id if scene_id else None,
                limit=limit
            )

            # 处理搜索结果
            memory_list = []
            for idx, result in enumerate(search_results):
                try:
                    # 从Mem0结果中提取信息
                    memory_content = result.get("memory", "") or result.get("content", "")
                    metadata = result.get("metadata", {})
                    score = result.get("score", 0.0)  # 相似度分数

                    # 解析对话内容
                    human_message = ""
                    ai_message = ""

                    if memory_content:
                        # 尝试解析对话格式
                        if "用户问:" in memory_content and "AI答:" in memory_content:
                            parts = memory_content.split("AI答:")
                            if len(parts) == 2:
                                human_message = parts[0].replace("用户问:", "").strip()
                                ai_message = parts[1].strip()
                        else:
                            # 如果不是标准格式，将整个内容作为记忆
                            human_message = memory_content
                            ai_message = ""

                    # 构建记忆字典
                    memory_dict = {
                        "id": f"vector_{idx}",  # 向量搜索结果的虚拟ID
                        "scene_id": metadata.get("scene_id", scene_id),
                        "agent_id": metadata.get("agent_id", agent_id),
                        "session_id": metadata.get("session_id", session_id),
                        "user_id": metadata.get("user_id", user_id),
                        "content": memory_content,
                        "human_message": human_message,
                        "ai_message": ai_message,
                        "similarity_score": round(score, 4),  # 相似度分数
                        "source": "vector_search",
                        "importance_level": 4 if score > 0.8 else 3,  # 根据相似度设置重要程度
                        "create_time": metadata.get("timestamp", datetime.now().isoformat()),
                        "metadata": metadata
                    }

                    # 根据memory_type进行筛选
                    if memory_type == "important" and memory_dict["similarity_score"] < 0.7:
                        continue
                    elif memory_type == "conversation" and not (human_message and ai_message):
                        continue
                    elif memory_type == "recent":
                        # 检查时间戳是否在最近7天内
                        try:
                            memory_time = datetime.fromisoformat(memory_dict["create_time"].replace('Z', '+00:00'))
                            from datetime import timedelta
                            if (datetime.now() - memory_time) > timedelta(days=7):
                                continue
                        except:
                            pass  # 如果时间解析失败，保留记录

                    memory_list.append(memory_dict)

                except Exception as parse_error:
                    print(f"解析搜索结果失败: {parse_error}")
                    continue

            # 按相似度分数排序（降序）
            memory_list.sort(key=lambda x: x.get("similarity_score", 0), reverse=True)

            # 构建返回结果
            result = {
                "success": True,
                "search_method": "vector_similarity",
                "query_keyword": keyword,
                "total_count": len(memory_list),
                "query_params": {
                    "scene_id": scene_id,
                    "agent_id": agent_id,
                    "session_id": session_id,
                    "user_id": user_id,
                    "keyword": keyword,
                    "tag": tag,
                    "limit": limit,
                    "memory_type": memory_type
                },
                "memories": memory_list,
                "similarity_threshold": 0.5,  # 相似度阈值
                "timestamp": datetime.now().isoformat(),
                "note": "结果基于向量相似度搜索，按相似度分数排序"
            }

            return json.dumps(result, ensure_ascii=False, indent=2)

        finally:
            db_session.close()

    except Exception as e:
        error_result = {
            "success": False,
            "error": f"向量搜索记忆历史失败: {str(e)}",
            "search_method": "vector_similarity",
            "timestamp": datetime.now().isoformat(),
            "suggestion": "请检查关键词并确保向量库正常运行"
        }
        return json.dumps(error_result, ensure_ascii=False)


@tool
def search_memory_by_content(
        query: str,
        scene_id: str = "",
        agent_id: str = "",
        user_id: str = "1",
        limit: int = 5
) -> str:
    """
    根据内容搜索相关记忆。

    参数:
        query: 必需，搜索查询内容
        scene_id: 可选，场景ID限制搜索范围
        agent_id: 可选，智能体ID限制搜索范围
        user_id: 可选，用户ID，默认为"1"
        limit: 可选，返回结果数量限制，默认5条
    返回:
        搜索到的相关记忆信息的JSON字符串。
    """
    try:
        if not query.strip():
            return json.dumps({"error": "搜索查询不能为空"}, ensure_ascii=False)

        # 获取数据库会话
        db_session = next(get_db_session())
        if not db_session:
            return json.dumps({"error": "无法连接数据库"}, ensure_ascii=False)

        try:
            # 尝试创建临时的记忆管理器进行搜索
            memory_manager = MemoryFactory.create_memory_system(
                memory_config={"type": "mem0"},
                db_session=db_session,
                scene_id=scene_id or "temp_search",
                agent_id=agent_id or "temp_search",
                session_id="temp_search_session",
                user_id=user_id
            )

            if memory_manager:
                # 执行同步搜索
                memories = memory_manager.search_relevant_memories(query, limit)

                result = {
                    "success": True,
                    "query": query,
                    "total_count": len(memories),
                    "search_params": {
                        "scene_id": scene_id,
                        "agent_id": agent_id,
                        "user_id": user_id,
                        "limit": limit
                    },
                    "memories": memories,
                    "timestamp": datetime.now().isoformat()
                }

                return json.dumps(result, ensure_ascii=False, indent=2)
            else:
                return json.dumps({"error": "无法创建记忆管理器"}, ensure_ascii=False)

        finally:
            db_session.close()

    except Exception as e:
        error_result = {
            "success": False,
            "error": f"搜索记忆失败: {str(e)}",
            "timestamp": datetime.now().isoformat()
        }
        return json.dumps(error_result, ensure_ascii=False)


class MemoryHistoryToolManager(ToolManager):
    """历史记忆工具管理器"""

    @property
    def category_name(self) -> str:
        return "获取历史记忆的工具"

    @property
    def category_code(self) -> str:
        return "memory_history"

    @property
    def category_description(self) -> str:
        return "提供记忆历史查询功能的工具集，包括根据各种条件获取和搜索历史记忆"

    def get_tool_configs(self) -> List[Dict[str, Any]]:
        # 获取记忆历史工具
        memory_tool = SingleToolConfig(
            tool_id=1,
            tool_name="get_memory_history",
            description="根据场景ID、Agent ID、会话ID、用户ID、关键词等条件获取历史记忆",
            is_enabled=True
        )

        memory_tool.add_parameter(
            param_name="scene_id",
            param_type="string",
            description="场景ID，用于筛选特定场景的记忆",
            default_value="",
            is_required=False
        )

        memory_tool.add_parameter(
            param_name="agent_id",
            param_type="string",
            description="智能体ID，用于筛选特定Agent的记忆",
            default_value="",
            is_required=False
        )

        memory_tool.add_parameter(
            param_name="session_id",
            param_type="string",
            description="会话ID，用于筛选特定对话的记忆",
            default_value="",
            is_required=False
        )

        memory_tool.add_parameter(
            param_name="user_id",
            param_type="string",
            description="用户ID",
            default_value="1",
            is_required=False
        )

        memory_tool.add_parameter(
            param_name="keyword",
            param_type="string",
            description="搜索关键词（必需），用于向量搜索相似内容的记忆",
            default_value="",
            is_required=True
        )

        memory_tool.add_parameter(
            param_name="tag",
            param_type="string",
            description="标签，用于筛选特定标签的记忆",
            default_value="",
            is_required=False
        )

        memory_tool.add_parameter(
            param_name="limit",
            param_type="integer",
            description="返回记忆数量限制",
            default_value="10",
            is_required=False
        )

        memory_tool.add_parameter(
            param_name="memory_type",
            param_type="string",
            description="记忆类型：all(所有)、conversation(对话)、important(重要)、recent(最近)",
            default_value="all",
            is_required=False
        )

        # 搜索记忆工具
        search_tool = SingleToolConfig(
            tool_id=2,
            tool_name="search_memory_by_content",
            description="根据内容搜索相关历史记忆",
            is_enabled=True
        )

        search_tool.add_parameter(
            param_name="query",
            param_type="string",
            description="搜索查询内容",
            default_value="",
            is_required=True
        )

        search_tool.add_parameter(
            param_name="scene_id",
            param_type="string",
            description="场景ID，限制搜索范围",
            default_value="",
            is_required=False
        )

        search_tool.add_parameter(
            param_name="agent_id",
            param_type="string",
            description="智能体ID，限制搜索范围",
            default_value="",
            is_required=False
        )

        search_tool.add_parameter(
            param_name="user_id",
            param_type="string",
            description="用户ID",
            default_value="1",
            is_required=False
        )

        search_tool.add_parameter(
            param_name="limit",
            param_type="integer",
            description="返回结果数量限制",
            default_value="5",
            is_required=False
        )

        return [memory_tool.to_dict(), search_tool.to_dict()]

    def get_langchain_tool(self):
        """返回工具列表"""
        return [get_memory_history, search_memory_by_content]


# 创建全局实例
memory_history_tool_manager = MemoryHistoryToolManager()
